#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
图像滑窗处理脚本 - u2net环境版本
功能：遍历指定文件夹中的所有图像，使用160x160的滑窗进行图像分割并保存
"""

import os
import time
from pathlib import Path

# 检查u2net环境
try:
    import cv2
    print("✓ OpenCV导入成功")
except ImportError:
    print("✗ 错误：无法导入OpenCV，请确保在u2net环境中运行")
    exit(1)

try:
    import numpy as np
    print("✓ NumPy导入成功")
except ImportError:
    print("✗ 错误：无法导入NumPy，请确保在u2net环境中运行")
    exit(1)

def create_sliding_windows(image, window_size=224, step_size=224):
    """对图像进行滑窗处理"""
    height, width = image.shape[:2]
    windows = []
    
    for y in range(0, height - window_size + 1, step_size):
        for x in range(0, width - window_size + 1, step_size):
            window = image[y:y + window_size, x:x + window_size]
            windows.append((window, x, y))
    
    return windows

def process_images(source_dir, target_dir, window_size=224):
    """处理源目录中的所有图像"""
    # 检查源目录是否存在
    if not os.path.exists(source_dir):
        print(f"✗ 错误：源目录不存在: {source_dir}")
        return
    
    # 创建目标目录
    Path(target_dir).mkdir(parents=True, exist_ok=True)
    print(f"✓ 目标目录已创建: {target_dir}")
    
    # 支持的图像格式
    image_extensions = {'.jpg', '.jpeg', '.png', '.bmp', '.tiff', '.tif'}
    image_files = []
    
    for file in os.listdir(source_dir):
        if Path(file).suffix.lower() in image_extensions:
            image_files.append(file)
    
    if not image_files:
        print("✗ 警告：源目录中没有找到图像文件")
        return
    
    print(f"✓ 找到 {len(image_files)} 个图像文件")
    
    total_windows = 0
    processed_count = 0
    
    for i, filename in enumerate(image_files, 1):
        print(f"\n处理图像 {i}/{len(image_files)}: {filename}")
        
        image_path = os.path.join(source_dir, filename)
        image = cv2.imread(image_path)
        
        if image is None:
            print(f"  ✗ 警告：无法读取图像 {filename}")
            continue
        
        print(f"  图像尺寸: {image.shape[1]}x{image.shape[0]}")
        
        name_without_ext = Path(filename).stem
        timestamp = int(time.time())
        windows = create_sliding_windows(image, window_size)
        
        if not windows:
            print(f"  ⚠ 图像尺寸过小，无法生成滑窗")
            continue
        
        for window, x, y in windows:
            output_filename = f"{name_without_ext}_{timestamp}_{x}_{y}.jpg"
            output_path = os.path.join(target_dir, output_filename)
            
            # 保存图像块
            success = cv2.imwrite(output_path, window)
            if success:
                total_windows += 1
            else:
                print(f"    ✗ 保存失败: {output_filename}")
        
        processed_count += 1
        print(f"  ✓ 生成了 {len(windows)} 个图像块")
    
    print(f"\n{'='*50}")
    print(f"处理完成！")
    print(f"成功处理图像: {processed_count}/{len(image_files)}")
    print(f"总共生成图像块: {total_windows}")
    print(f"保存位置: {target_dir}")

def main():
    """主函数"""
    default_source = "/Volumes/data1/JH/projects/JLD/JLD_quyang/0828/RC250826012/06P/info.imgs"
    default_target = "/Volumes/data1/JH/dataset/06P"
    
    print("="*60)
    print("图像滑窗处理工具 - u2net环境版本")
    print("="*60)
    print(f"源目录: {default_source}")
    print(f"目标目录: {default_target}")
    print(f"滑窗大小: 224x224像素")
    print(f"滑窗步长: 224像素（无重叠）")
    print("-" * 60)
    
    # 检查环境
    print("环境检查:")
    
    start_time = time.time()
    process_images(default_source, default_target, 224)
    end_time = time.time()
    
    print(f"\n总耗时: {end_time - start_time:.2f} 秒")
    print("="*60)

if __name__ == "__main__":
    main()
